"""心理咨询ReAct Agent - 集成向量检索的智能对话代理"""

from typing import Dict, Any, List, Optional, AsyncGenerator
from datetime import datetime
import logging
import json
import asyncio

from .base import BaseAgent
from ..tools.base import tool_registry
from ..tools.intent_analysis import IntentAnalysisTool
from ..tools.safety_check import SafetyCheckTool
from ..tools.knowledge_retrieval import KnowledgeRetrievalTool
from ..tools.context_retrieval import ContextRetrievalTool, ConversationMemoryManager
from ..services.conversation_service import ConversationService
from ...models.schemas.chat import ReActStep, ReActResponse

logger = logging.getLogger(__name__)


class PsychologicalAgent(BaseAgent):
    """心理咨询ReAct Agent - 实现完整的推理-行动-观察循环"""
    
    def __init__(self, db_session):
        super().__init__(
            name="psychological_counseling_agent",
            description="专业的心理咨询ReAct Agent，提供智能化的心理健康支持"
        )
        
        # 初始化服务
        self.conversation_service = ConversationService(db_session)
        self.memory_manager = ConversationMemoryManager()
        
        # 注册工具
        self._register_tools()
        
        # ReAct配置
        self.max_steps = 10
        self.thinking_timeout = 30  # 秒
        
    def _register_tools(self):
        """注册ReAct工具"""
        tools = [
            IntentAnalysisTool(),
            SafetyCheckTool(),
            KnowledgeRetrievalTool(),
            ContextRetrievalTool()
        ]
        
        for tool in tools:
            tool_registry.register(tool)
            
        logger.info(f"已注册 {len(tools)} 个ReAct工具")
    
    async def process_message(self, user_input: str, conversation_id: str,
                            user_id: Optional[str] = None, 
                            stream: bool = True) -> AsyncGenerator[Dict[str, Any], None]:
        """处理用户消息 - 主要入口点"""
        start_time = datetime.now()
        
        try:
            # 1. 初始化对话上下文
            conversation = await self._ensure_conversation_exists(conversation_id, user_id)
            
            # 2. 保存用户消息
            user_message = self.conversation_service.add_message(
                conversation_id=conversation_id,
                role="user",
                content=user_input,
                metadata={"agent_version": "v1.0", "processing_start": start_time.isoformat()}
            )
            
            # 3. 添加到向量记忆
            await self._add_message_to_memory(user_message)
            
            # 4. 执行ReAct推理循环
            react_steps = []
            final_response = ""
            
            if stream:
                async for step_data in self._execute_react_loop_stream(
                    user_input, conversation_id, user_id, react_steps
                ):
                    yield step_data
                    
                    if step_data.get("type") == "final_response":
                        final_response = step_data.get("content", "")
            else:
                final_response = await self._execute_react_loop(
                    user_input, conversation_id, user_id, react_steps
                )
                
                # 返回完整结果
                yield {
                    "type": "complete_response",
                    "content": final_response,
                    "steps": react_steps,
                    "conversation_id": conversation_id,
                    "timestamp": datetime.now().isoformat()
                }
            
            # 5. 保存助手回复
            assistant_message = self.conversation_service.add_message(
                conversation_id=conversation_id,
                role="assistant", 
                content=final_response,
                metadata={
                    "react_steps": len(react_steps),
                    "processing_time": (datetime.now() - start_time).total_seconds(),
                    "agent_version": "v1.0"
                }
            )
            
            # 6. 添加助手回复到向量记忆
            await self._add_message_to_memory(assistant_message)
            
        except Exception as e:
            logger.error(f"处理消息失败: {e}")
            yield {
                "type": "error",
                "content": f"处理您的消息时遇到了问题: {str(e)}",
                "timestamp": datetime.now().isoformat()
            }
    
    async def _execute_react_loop_stream(self, user_input: str, conversation_id: str,
                                       user_id: Optional[str], react_steps: List[ReActStep]) -> AsyncGenerator[Dict[str, Any], None]:
        """流式执行ReAct循环"""
        step_count = 0
        
        # 步骤1: 思考阶段
        yield {"type": "thinking", "content": "正在分析您的问题...", "step": step_count + 1}
        
        # 获取增强上下文
        context_data = await self.memory_manager.get_enhanced_context(
            user_input, conversation_id, user_id
        )
        
        # 步骤2: 意图分析
        step_count += 1
        yield {"type": "action", "content": "分析用户意图和情绪状态", "step": step_count}
        
        intent_tool = tool_registry.get_tool("analyze_user_intent")
        intent_result = await intent_tool.safe_execute({
            "user_input": user_input,
            "chat_history": json.dumps(context_data.get("current_context", []))
        })
        
        react_steps.append(ReActStep(
            step=step_count,
            type="action",
            content="意图分析",
            tool_name="analyze_user_intent",
            tool_input={"user_input": user_input},
            tool_output=json.dumps(intent_result.result) if intent_result.success else intent_result.error,
            timestamp=datetime.now().isoformat()
        ))
        
        yield {"type": "observation", "content": f"识别意图: {intent_result.intent if intent_result.success else '分析失败'}", "step": step_count}
        
        # 步骤3: 安全检查
        step_count += 1
        yield {"type": "action", "content": "进行安全风险评估", "step": step_count}
        
        safety_tool = tool_registry.get_tool("check_user_safety")
        safety_result = await safety_tool.safe_execute({
            "user_input": user_input,
            "chat_history": json.dumps(context_data.get("current_context", []))
        })
        
        react_steps.append(ReActStep(
            step=step_count,
            type="action", 
            content="安全检查",
            tool_name="check_user_safety",
            tool_input={"user_input": user_input},
            tool_output=json.dumps(safety_result.result) if safety_result.success else safety_result.error,
            timestamp=datetime.now().isoformat()
        ))
        
        risk_level = safety_result.risk_level if safety_result.success else "unknown"
        yield {"type": "observation", "content": f"风险等级: {risk_level}", "step": step_count}
        
        # 步骤4: 知识检索
        step_count += 1
        yield {"type": "action", "content": "检索相关心理健康知识", "step": step_count}
        
        knowledge_tool = tool_registry.get_tool("retrieve_knowledge_documents")
        knowledge_result = await knowledge_tool.safe_execute({
            "user_input": user_input,
            "intent": intent_result.intent if intent_result.success else "general_consultation"
        })
        
        react_steps.append(ReActStep(
            step=step_count,
            type="action",
            content="知识检索",
            tool_name="retrieve_knowledge_documents", 
            tool_input={"user_input": user_input},
            tool_output=json.dumps(knowledge_result.result) if knowledge_result.success else knowledge_result.error,
            timestamp=datetime.now().isoformat()
        ))
        
        docs_found = len(knowledge_result.documents) if knowledge_result.success else 0
        yield {"type": "observation", "content": f"找到 {docs_found} 个相关文档", "step": step_count}
        
        # 步骤5: 生成回复
        step_count += 1
        yield {"type": "thinking", "content": "基于分析结果生成专业回复...", "step": step_count}
        
        final_response = await self._generate_final_response(
            user_input=user_input,
            intent_result=intent_result,
            safety_result=safety_result,
            knowledge_result=knowledge_result,
            context_data=context_data
        )
        
        react_steps.append(ReActStep(
            step=step_count,
            type="thinking",
            content="生成最终回复",
            tool_name=None,
            tool_input=None,
            tool_output=final_response,
            timestamp=datetime.now().isoformat()
        ))
        
        yield {"type": "final_response", "content": final_response, "step": step_count}
    
    async def _execute_react_loop(self, user_input: str, conversation_id: str,
                                user_id: Optional[str], react_steps: List[ReActStep]) -> str:
        """非流式执行ReAct循环"""
        # 获取增强上下文
        context_data = await self.memory_manager.get_enhanced_context(
            user_input, conversation_id, user_id
        )
        
        # 执行所有工具
        intent_tool = tool_registry.get_tool("analyze_user_intent")
        safety_tool = tool_registry.get_tool("check_user_safety")
        knowledge_tool = tool_registry.get_tool("retrieve_knowledge_documents")
        
        # 并行执行工具（除了依赖关系的）
        intent_result, safety_result = await asyncio.gather(
            intent_tool.safe_execute({
                "user_input": user_input,
                "chat_history": json.dumps(context_data.get("current_context", []))
            }),
            safety_tool.safe_execute({
                "user_input": user_input,
                "chat_history": json.dumps(context_data.get("current_context", []))
            })
        )
        
        # 基于意图结果执行知识检索
        knowledge_result = await knowledge_tool.safe_execute({
            "user_input": user_input,
            "intent": intent_result.intent if intent_result.success else "general_consultation"
        })
        
        # 记录步骤
        react_steps.extend([
            ReActStep(step=1, type="action", content="意图分析", tool_name="analyze_user_intent", 
                     tool_input={"user_input": user_input}, 
                     tool_output=json.dumps(intent_result.result) if intent_result.success else intent_result.error,
                     timestamp=datetime.now().isoformat()),
            ReActStep(step=2, type="action", content="安全检查", tool_name="check_user_safety",
                     tool_input={"user_input": user_input},
                     tool_output=json.dumps(safety_result.result) if safety_result.success else safety_result.error,
                     timestamp=datetime.now().isoformat()),
            ReActStep(step=3, type="action", content="知识检索", tool_name="retrieve_knowledge_documents",
                     tool_input={"user_input": user_input},
                     tool_output=json.dumps(knowledge_result.result) if knowledge_result.success else knowledge_result.error,
                     timestamp=datetime.now().isoformat())
        ])
        
        # 生成最终回复
        final_response = await self._generate_final_response(
            user_input=user_input,
            intent_result=intent_result,
            safety_result=safety_result,
            knowledge_result=knowledge_result,
            context_data=context_data
        )
        
        return final_response
    
    async def _generate_final_response(self, user_input: str, intent_result, safety_result,
                                     knowledge_result, context_data: Dict[str, Any]) -> str:
        """生成最终回复"""
        try:
            # 如果是高风险情况，优先使用安全回复
            if safety_result.success and safety_result.risk_level == "high":
                return safety_result.response
            
            # 构建回复上下文
            response_context = {
                "user_input": user_input,
                "intent": intent_result.intent if intent_result.success else "general_consultation",
                "emotion": intent_result.emotion if intent_result.success else "neutral",
                "risk_level": safety_result.risk_level if safety_result.success else "low",
                "knowledge_docs": knowledge_result.documents if knowledge_result.success else [],
                "conversation_context": context_data.get("context_summary", ""),
                "similar_conversations": context_data.get("similar_conversations", [])
            }
            
            # 生成个性化回复
            response = await self._create_personalized_response(response_context)
            
            return response
            
        except Exception as e:
            logger.error(f"生成最终回复失败: {e}")
            return "我理解您的困扰，虽然在处理您的问题时遇到了一些技术问题，但我想让您知道，您的感受是重要的。如果您愿意，可以告诉我更多关于您的情况，我会尽力帮助您。"
    
    async def _create_personalized_response(self, context: Dict[str, Any]) -> str:
        """创建个性化回复"""
        # 这里应该调用LLM生成回复
        # 暂时使用基于规则的回复生成
        
        intent = context.get("intent", "general_consultation")
        emotion = context.get("emotion", "neutral")
        risk_level = context.get("risk_level", "low")
        
        # 基础回复模板
        if intent == "anxiety":
            base_response = "我理解您现在感到焦虑。焦虑是一种常见的情绪反应，让我们一起探讨一些应对策略。"
        elif intent == "depression":
            base_response = "我能感受到您现在的沮丧情绪。抑郁的感受确实很难承受，但请记住您并不孤单。"
        elif intent == "stress":
            base_response = "工作和生活的压力确实会让人感到疲惫。让我们一起找到一些有效的压力管理方法。"
        elif intent == "sleep":
            base_response = "睡眠问题会影响我们的整体健康。让我为您提供一些改善睡眠质量的建议。"
        else:
            base_response = "感谢您与我分享您的感受。我在这里倾听并支持您。"
        
        # 添加个性化建议
        if context.get("knowledge_docs"):
            base_response += "\n\n基于相关的心理健康知识，我建议您可以尝试以下方法："
            # 这里可以从知识文档中提取建议
        
        # 添加上下文相关的支持
        if context.get("similar_conversations"):
            base_response += "\n\n从您之前的对话中，我注意到您一直在努力应对这些挑战，这显示了您的坚韧和勇气。"
        
        return base_response
    
    async def _ensure_conversation_exists(self, conversation_id: str, user_id: Optional[str]):
        """确保对话存在"""
        return self.conversation_service.get_or_create_conversation(
            user_id=user_id,
            conversation_id=conversation_id,
            title="心理咨询对话"
        )
    
    async def _add_message_to_memory(self, message):
        """添加消息到向量记忆"""
        try:
            message_data = {
                "message_id": message.message_id,
                "conversation_id": message.conversation_id,
                "role": message.role,
                "content": message.content,
                "created_at": message.created_at.isoformat() if message.created_at else None,
                "metadata": message.get_metadata()
            }
            
            await self.memory_manager.add_message_to_memory(message_data)
            
        except Exception as e:
            logger.error(f"添加消息到向量记忆失败: {e}")
    
    def get_agent_info(self) -> Dict[str, Any]:
        """获取Agent信息"""
        return {
            "name": self.name,
            "description": self.description,
            "version": "1.0.0",
            "capabilities": [
                "意图识别和情绪分析",
                "安全风险评估",
                "心理健康知识检索",
                "对话上下文记忆",
                "个性化回复生成",
                "流式对话处理"
            ],
            "tools": list(tool_registry.get_tool_names()),
            "max_steps": self.max_steps,
            "supports_streaming": True
        }
